#include <iostream>
#include <string>
#include <vector>

// 前向声明命令函数
namespace aicompiler {
    int cmdIngest(const std::string& input_file);
    int cmdOpt(const std::string& input_file, const std::string& output_file);
    int cmdLower(const std::string& input_file, const std::string& target, 
                 const std::string& output_file);
    int cmdCodegen(const std::string& input_file, const std::string& output_file);
    int cmdRun(const std::string& executable);
}

void printUsage() {
    std::cout << "AIcompiler 0.1.0 - Universal compiler middle+back-end\n\n"
              << "Usage: aicompiler <command> [options]\n\n"
              << "Commands:\n"
              << "  ingest <file.mlir>              Parse and verify MLIR input\n"
              << "  opt <input> -o <output>         Apply optimization passes\n"
              << "  lower <input> --target=<tgt> -o <output>\n"
              << "                                  Lower to target backend (cpu|spirv|nvptx)\n"
              << "  codegen <input> -o <output>     Generate executable\n"
              << "  run <executable>                Execute compiled program\n\n"
              << "Examples:\n"
              << "  aicompiler ingest examples/vector_add.mlir\n"
              << "  aicompiler opt input.mlir -o optimized.mlir\n"
              << "  aicompiler lower optimized.mlir --target=cpu -o output.ll\n";
}

int main(int argc, char** argv) {
    if (argc <= 1) {
        printUsage();
        return 0;
    }
    
    std::string cmd = argv[1];
    std::vector<std::string> args(argv + 2, argv + argc);
    
    try {
        if (cmd == "ingest") {
            if (args.empty()) {
                std::cerr << "Error: ingest requires input file\n";
                return 1;
            }
            return aicompiler::cmdIngest(args[0]);
        }
        else if (cmd == "opt") {
            if (args.size() < 3 || args[1] != "-o") {
                std::cerr << "Error: opt requires <input> -o <output>\n";
                return 1;
            }
            return aicompiler::cmdOpt(args[0], args[2]);
        }
        else if (cmd == "lower") {
            std::string input, target = "cpu", output;
            for (size_t i = 0; i < args.size(); ++i) {
                if (args[i] == "-o" && i + 1 < args.size()) {
                    output = args[++i];
                } else if (args[i].find("--target=") == 0) {
                    target = args[i].substr(9);
                } else if (input.empty()) {
                    input = args[i];
                }
            }
            if (input.empty() || output.empty()) {
                std::cerr << "Error: lower requires <input> --target=<tgt> -o <output>\n";
                return 1;
            }
            return aicompiler::cmdLower(input, target, output);
        }
        else if (cmd == "codegen") {
            if (args.size() < 3 || args[1] != "-o") {
                std::cerr << "Error: codegen requires <input> -o <output>\n";
                return 1;
            }
            return aicompiler::cmdCodegen(args[0], args[2]);
        }
        else if (cmd == "run") {
            if (args.empty()) {
                std::cerr << "Error: run requires executable path\n";
                return 1;
            }
            return aicompiler::cmdRun(args[0]);
        }
        else if (cmd == "--help" || cmd == "-h") {
            printUsage();
            return 0;
        }
        else {
            std::cerr << "Unknown command: " << cmd << "\n";
            printUsage();
            return 1;
        }
    }
    catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}
